package net.xo68.destiny.core.rocksdb;


import org.rocksdb.*;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.File;
import java.io.IOException;
import java.util.*;
import java.util.concurrent.locks.Lock;
import java.util.concurrent.locks.ReadWriteLock;
import java.util.concurrent.locks.ReentrantReadWriteLock;

/**
 * rocksdb 存储
 * @author wuxie
 * @version 2020-11-4
 */
public class RocksDBStorage implements AutoCloseable {

    private static final Logger log = LoggerFactory.getLogger(RocksDBStorage.class);

    static {
        RocksDB.loadLibrary();
    }

    private final String path;
    private final boolean sync;
    private final boolean openStatistics;
    private RocksDB db;
    private DBOptions dbOptions;
    private WriteOptions writeOptions;
    private final List<ColumnFamilyOptions> cfOptions= new ArrayList<>();
    private ColumnFamilyHandle defaultHandle;
    private ColumnFamilyHandle metadataHandle;
    private ReadOptions totalOrderReadOptions;
    private DebugStatistics statistics;

    private final ReadWriteLock readWriteLock = new ReentrantReadWriteLock();
    private final Lock readLock= this.readWriteLock.readLock();
    private final Lock writeLock= this.readWriteLock.writeLock();

    private RocksSerializer<String> stringKeySerializer;
    private RocksSerializer<Long> longKeySerializer;
    private RocksSerializer valueSerializer;

    public RocksDBStorage(final String path, boolean sync,boolean openStatistics) {
        this.path = path;
        this.sync = sync;
        this.openStatistics=openStatistics;
    }

    public void setStringKeySerializer(RocksSerializer stringKeySerializer) {
        this.stringKeySerializer = stringKeySerializer;
    }

    public void setLongKeySerializer(RocksSerializer longKeySerializer) {
        this.longKeySerializer = longKeySerializer;
    }

    public void setValueSerializer(RocksSerializer valueSerializer) {
        this.valueSerializer = valueSerializer;
    }

    public static DBOptions createDBOptions() {
        return RocksDBOptionFactory.getRocksDBOptions(RocksDBStorage.class);
    }
    public static ColumnFamilyOptions createColumnFamilyOptions() {
        final BlockBasedTableConfig tConfig = RocksDBOptionFactory
                .getRocksDBTableFormatConfig(RocksDBStorage.class);
        return RocksDBOptionFactory.getRocksDBColumnFamilyOptions(RocksDBStorage.class) //
                .useFixedLengthPrefixExtractor(8) //
                .setTableFormatConfig(tConfig) //
                .setMergeOperator(new StringAppendOperator());
    }

    public boolean init(){
        this.writeLock.lock();
        try {
            if (this.db != null) {
                log.warn("RocksDBLogStorage init() already.");
                return true;
            }
            if(this.stringKeySerializer==null){
                this.stringKeySerializer=new StringRocksSerializer();
            }
            if(this.longKeySerializer==null){
                this.longKeySerializer=new LongRocksSerializer();
            }
            if(this.valueSerializer==null){
                this.valueSerializer=new GenericJackson2JsonRocksValueSerializer(Object.class);
            }
            this.dbOptions = createDBOptions();
            if (this.openStatistics) {
                this.statistics = new DebugStatistics();
                this.dbOptions.setStatistics(this.statistics);
            }
            this.writeOptions = new WriteOptions();
            this.writeOptions.setSync(this.sync);
            this.totalOrderReadOptions = new ReadOptions();
            this.totalOrderReadOptions.setTotalOrderSeek(true);
            return initAndLoad();
        } catch (final RocksDBException e) {
            log.error("Fail to init RocksDBLogStorage, path={}.", this.path, e);
            return false;
        } finally {
            this.writeLock.unlock();
        }
    }

    private boolean initAndLoad() throws RocksDBException{
        final List<ColumnFamilyDescriptor> columnFamilyDescriptors = new ArrayList<>();
        final ColumnFamilyOptions cfOption = createColumnFamilyOptions();
        this.cfOptions.add(cfOption);
        // Column family to store configuration log entry.
        columnFamilyDescriptors.add(new ColumnFamilyDescriptor("Metadata".getBytes(), cfOption));
        // Default column family to store user data log entry.
        columnFamilyDescriptors.add(new ColumnFamilyDescriptor(RocksDB.DEFAULT_COLUMN_FAMILY, cfOption));
        openDB(columnFamilyDescriptors);
        //todo load metadata
        return true;
    }

    private void openDB(final List<ColumnFamilyDescriptor> columnFamilyDescriptors) throws RocksDBException {
        final List<ColumnFamilyHandle> columnFamilyHandles = new ArrayList<>();

        final File dir = new File(this.path);
        //不是目录
        if (dir.exists() && !dir.isDirectory()) {
            throw new IllegalStateException("Invalid log path, it's a regular file: " + this.path);
        }
        this.db = RocksDB.open(this.dbOptions, this.path, columnFamilyDescriptors, columnFamilyHandles);

        assert (columnFamilyHandles.size() == 2);
        this.metadataHandle = columnFamilyHandles.get(0);
        this.defaultHandle = columnFamilyHandles.get(1);
    }

    private void checkState() {
        if(this.db==null){
            throw new NullPointerException("DB not initialized or destroyed");
        }
    }

    public byte[] get(final String key){
        final byte[] keyBytes=this.stringKeySerializer.serialize(key);
        return get(keyBytes);
    }
    public byte[] get(final Long key){
        final byte[] keyBytes=this.longKeySerializer.serialize(key);
        return get(keyBytes);
    }
    public <T> T getObject(final String key){
        final byte[] valueBytes=this.get(key);
        return (T) this.valueSerializer.deserialize(valueBytes);
    }
    public <T> T getObject(final Long key){
        final byte[] valueBytes=this.get(key);
        return (T) this.valueSerializer.deserialize(valueBytes);
    }
    public byte[] get(final byte[] key){
        this.readLock.lock();
        try {
            return this.db.get(this.defaultHandle, key);
        } catch (final RocksDBException e) {
            log.error("Fail to get log entry at index {}.",  e);
        } finally {
            this.readLock.unlock();
        }
        return null;
    }

    public <T> Map<String,T> limitString(long count){
        return limitString(0,count);
    }
    public <T> Map<String,T> limitString(long offset,long count){
        RocksIterator iter = this.db.newIterator(this.defaultHandle,this.totalOrderReadOptions);
        int pos=0;
        Map<String,T> resultMap= new HashMap<>();
        for (iter.seekToFirst();iter.isValid();iter.next()){
            if(pos >= offset && pos < offset+count+1){
                resultMap.put(this.stringKeySerializer.deserialize(iter.key()),
                        (T) this.valueSerializer.deserialize(iter.value()));
            }
            if(pos >= offset+count+1){
                break;
            }
            pos++;
        }
        return resultMap;
    }
    public <T> Map<Long,T> limitLong(long count){
        return limitLong(0L,count);
    }
    public <T> Map<Long,T> limitLong(long offset,long count){
        RocksIterator iter = this.db.newIterator(this.defaultHandle,this.totalOrderReadOptions);
        int pos=0;
        Map<Long,T> resultMap= new LinkedHashMap<>();
        for (iter.seekToFirst();iter.isValid();iter.next()){
            if((pos >= offset && count==-1L) || (pos >= offset && pos < offset+count)){
                final long key=this.longKeySerializer.deserialize(iter.key());
                //log.info("iter key:{}",key);
                resultMap.put(key,
                        (T) this.valueSerializer.deserialize(iter.value()));
            }
            if(count!=-1L && pos >= offset+count){
                break;
            }
            pos++;
        }
        return resultMap;
    }

    public Map<byte[],byte[]> limit(long count){
        return limit(0L,count);
    }
    public Map<byte[],byte[]> limit(long offset,long count){
        RocksIterator iter = this.db.newIterator(this.defaultHandle,this.totalOrderReadOptions);
        int pos=0;
        Map<byte[],byte[]> resultMap= new LinkedHashMap<>();
        for (iter.seekToFirst();iter.isValid();iter.next()){
            if((pos >= offset && count==-1L) || (pos >= offset && pos < offset+count)){
                resultMap.put(iter.key(),iter.value());
            }
            if(count!=-1L && pos >= offset+count){
                break;
            }
            pos++;
        }
        return resultMap;
    }

    public <T> boolean put(String key,T t){
        final byte[] keyBytes=this.stringKeySerializer.serialize(key);
        final byte[] valBytes=this.valueSerializer.serialize(t);
        return put(keyBytes,valBytes);
    }
    public <T> boolean put(Long key,T t){
        final byte[] keyBytes=this.longKeySerializer.serialize(key);
        final byte[] valBytes=this.valueSerializer.serialize(t);
        return put(keyBytes,valBytes);
    }
    public boolean put(final byte[] key,final byte[] val){
        this.readLock.lock();
        try {
            if (this.db == null) {
                log.warn("DB not initialized or destroyed.");
                return false;
            }
            this.db.put(this.defaultHandle, this.writeOptions, key, val);
            return true;
        } catch (final RocksDBException e) {
            log.error("Fail to append entry.", e);
            return false;
        } finally {
            this.readLock.unlock();
        }
    }



    public <T> int putStringAll(final Map<String,T> itemMap){
        if(itemMap==null || itemMap.isEmpty()){
            return 0;
        }
        final int itemMapCount=itemMap.size();
        final boolean ret= executeBatch(batch -> {
            for (Map.Entry<String,T> item:itemMap.entrySet()){
                batch.put(this.defaultHandle,this.stringKeySerializer.serialize(item.getKey()),
                        this.valueSerializer.serialize(item.getValue()));
            }
        });
        if(ret){
            return itemMapCount;
        }else {
            return 0;
        }
    }
    public <T> int putLongAll(final Map<Long,T> itemMap){
        if(itemMap==null || itemMap.isEmpty()){
            return 0;
        }
        final int itemMapCount=itemMap.size();
        final boolean ret= executeBatch(batch -> {
            for (Map.Entry<Long,T> item:itemMap.entrySet()){
                batch.put(this.defaultHandle,this.longKeySerializer.serialize(item.getKey()),
                        this.valueSerializer.serialize(item.getValue()));
            }
        });
        if(ret){
            return itemMapCount;
        }else {
            return 0;
        }
    }


    public int putAll(final Map<byte[],byte[]> itemMap){
        if(itemMap==null || itemMap.isEmpty()){
            return 0;
        }
        final int itemMapCount=itemMap.size();
        final boolean ret= executeBatch(batch -> {
            for (Map.Entry<byte[],byte[]> item:itemMap.entrySet()){
                batch.put(this.defaultHandle, item.getKey(), item.getValue());
            }
        });
        if(ret){
            return itemMapCount;
        }else {
            return 0;
        }
    }

    public boolean remove(final String key){
        final byte[] keyBytes=this.stringKeySerializer.serialize(key);
        return remove(keyBytes);
    }
    public boolean remove(final Long key){
        final byte[] keyBytes=this.longKeySerializer.serialize(key);
        return remove(keyBytes);
    }
    public boolean remove(final byte[] key){
        this.readLock.lock();
        try {
            if (this.db == null) {
                log.warn("DB not initialized or destroyed.");
                return false;
            }
            this.db.delete(this.defaultHandle, key);
            return true;
        } catch (final RocksDBException e) {
            log.error("Fail to append entry.", e);
            return false;
        } finally {
            this.readLock.unlock();
        }
    }

    public boolean exist(String key){
        return exist(this.stringKeySerializer.serialize(key));
    }
    public boolean exist(Long key){
        return exist(this.longKeySerializer.serialize(key));
    }
    public boolean exist(final byte[] key){
        final Holder<byte[]> holder=new Holder<>();
        return this.db.keyMayExist(this.defaultHandle,key,holder);
    }

    /**
     * Execute write batch template.
     *
     * @param template write batch template
     */
    private boolean executeBatch(final WriteBatchTemplate template) {
        this.readLock.lock();
        if (this.db == null) {
            log.warn("DB not initialized or destroyed.");
            this.readLock.unlock();
            return false;
        }
        try (final WriteBatch batch = new WriteBatch()) {
            template.execute(batch);
            this.db.write(this.writeOptions, batch);
        } catch (final RocksDBException e) {
            log.error("Execute batch failed with rocksdb exception.", e);
            return false;
        } catch (final IOException e) {
            log.error("Execute batch failed with io exception.", e);
            return false;
        } catch (final InterruptedException e) {
            log.error("Execute batch failed with interrupt.", e);
            Thread.currentThread().interrupt();
            return false;
        } finally {
            this.readLock.unlock();
        }
        return true;
    }

    public void destroyDB(){
        final Options opt = new Options();
        try {
            RocksDB.destroyDB(this.path, opt);
        } catch (RocksDBException e) {
            e.printStackTrace();
        }
    }

    @Override
    public void close() throws Exception {
        this.writeLock.lock();
        try {
            // The shutdown order is matter.
            // 1. close column family handles
            this.metadataHandle.close();
            this.defaultHandle.close();
            this.db.close();
            // 2. close column family options.
            for (final ColumnFamilyOptions opt : this.cfOptions) {
                opt.close();
            }
            // 3. close options
            this.dbOptions.close();
            if (this.statistics != null) {
                this.statistics.close();
            }
            this.writeOptions.close();
            this.totalOrderReadOptions.close();
            // 4. help gc.
            this.cfOptions.clear();
            this.dbOptions = null;
            this.statistics = null;
            this.writeOptions = null;
            this.totalOrderReadOptions = null;
            this.defaultHandle = null;
            this.metadataHandle = null;
            this.db = null;
            log.info("DB destroyed, the db path is: {}.", this.path);
        } finally {
            this.writeLock.unlock();
        }
    }

    /**
     * Write batch template.
     *
     * @author boyan (boyan@alibaba-inc.com)
     *
     * 2017-Nov-08 11:19:22 AM
     */
    public static interface WriteBatchTemplate {
        void execute(WriteBatch batch) throws RocksDBException, IOException, InterruptedException;
    }
}
